import numpy as np
import time
import matplotlib.pyplot as plt
from qp_taichi import QPSolver
import taichi as ti

def generate_random_qp(n, m_ineq, m_eq, seed=42):
    """
    生成随机QP问题
    
    :param n: 变量数量
    :param m_ineq: 不等式约束数量
    :param m_eq: 等式约束数量
    :param seed: 随机种子
    :return: P, q, G, h, A, b - QP问题参数
    """
    np.random.seed(seed)
    
    # 生成随机半正定矩阵 P
    X = np.random.randn(n, n)
    P = X.T @ X + 0.1 * np.eye(n)  # 确保正定性
    
    # 生成随机向量 q
    q = np.random.randn(n)
    
    # 生成随机不等式约束
    if m_ineq > 0:
        G = np.random.randn(m_ineq, n)
        h = np.abs(np.random.randn(m_ineq)) * 10  # 确保约束有解
    else:
        G, h = None, None
    
    # 生成随机等式约束（确保满秩）
    if m_eq > 0:
        # 确保 m_eq <= n，否则系统可能无解
        m_eq_actual = min(m_eq, n-1)  # 保留至少一个自由度
        A_tmp = np.random.randn(m_eq_actual, n)
        # 通过SVD确保满秩
        u, s, vh = np.linalg.svd(A_tmp, full_matrices=False)
        A = u @ vh
        # 确保系统有解
        x_random = np.random.randn(n)
        b = A @ x_random
    else:
        A, b = None, None
    
    return P, q, G, h, A, b

def benchmark_qp_solver(sizes, num_runs=3, use_gpu=True):
    """
    对不同规模的QP问题进行基准测试
    
    :param sizes: 变量数量的列表
    :param num_runs: 每个规模运行的次数
    :param use_gpu: 是否使用GPU
    :return: 平均求解时间列表
    """
    arch = ti.gpu if use_gpu else ti.cpu
    solver = QPSolver(arch=arch)
    
    avg_times = []
    
    for n in sizes:
        print(f"\n测试规模: n = {n}")
        
        # 定义约束数量 (通常不等式约束多于等式约束)
        m_ineq = n
        m_eq = n // 4
        
        times = []
        
        for run in range(num_runs):
            # 生成随机问题
            P, q, G, h, A, b = generate_random_qp(n, m_ineq, m_eq, seed=42+run)
            
            # 测量求解时间
            t_start = time.time()
            result = solver.solve(P, q, G, h, A, b, 
                                  max_iter=500,
                                  abs_tol=1e-3, 
                                  rel_tol=1e-2,
                                  verbose=False)
            t_end = time.time()
            
            solve_time = t_end - t_start
            times.append(solve_time)
            
            # 打印本次运行结果
            print(f"  运行 {run+1}/{num_runs}, 迭代: {result.iterations}, " 
                  f"状态: {result.status}, 时间: {solve_time:.3f}秒")
        
        # 计算平均求解时间
        avg_time = np.mean(times)
        avg_times.append(avg_time)
        print(f"  平均求解时间: {avg_time:.3f} 秒")
    
    return avg_times

def plot_benchmark_results(sizes, gpu_times, cpu_times):
    """
    绘制基准测试结果
    
    :param sizes: 问题规模列表
    :param gpu_times: GPU求解时间列表
    :param cpu_times: CPU求解时间列表
    """
    plt.figure(figsize=(10, 6))
    
    plt.plot(sizes, gpu_times, 'o-', label='GPU')
    plt.plot(sizes, cpu_times, 's-', label='CPU')
    
    plt.xlabel('问题规模 (变量数量)')
    plt.ylabel('求解时间 (秒)')
    plt.title('QP-Taichi: GPU vs CPU 性能比较')
    plt.grid(True)
    plt.legend()
    
    # 计算加速比
    speedups = [cpu/gpu for cpu, gpu in zip(cpu_times, gpu_times)]
    
    # 添加加速比文本标签
    for i, (size, speedup) in enumerate(zip(sizes, speedups)):
        plt.annotate(f'{speedup:.1f}x', 
                     xy=(size, gpu_times[i]),
                     xytext=(0, 10),
                     textcoords='offset points',
                     ha='center')
    
    plt.tight_layout()
    plt.savefig('benchmark_results.png')
    plt.close()

def main():
    """运行基准测试并比较GPU和CPU性能"""
    print("QP-Taichi 性能基准测试: GPU vs CPU")
    
    # 定义测试的问题规模
    sizes = [10, 50, 100, 500, 1000, 2000, 5000]
    
    # GPU基准测试
    try:
        print("\n运行GPU基准测试...")
        gpu_times = benchmark_qp_solver(sizes, use_gpu=True)
    except Exception as e:
        print(f"GPU测试失败: {e}")
        gpu_times = [float('inf')] * len(sizes)
    
    # CPU基准测试
    print("\n运行CPU基准测试...")
    cpu_times = benchmark_qp_solver(sizes, use_gpu=False)
    
    # 绘制结果
    try:
        print("\n生成性能比较图...")
        plot_benchmark_results(sizes, gpu_times, cpu_times)
        print("结果已保存到 benchmark_results.png")
    except Exception as e:
        print(f"绘图失败: {e}")
    
    # 打印性能比较表格
    print("\n性能比较结果:")
    print("-" * 60)
    print("问题规模  |  GPU时间(秒)  |  CPU时间(秒)  |  加速比(CPU/GPU)")
    print("-" * 60)
    
    for i, size in enumerate(sizes):
        gpu_time = gpu_times[i]
        cpu_time = cpu_times[i]
        
        if gpu_time > 0 and gpu_time < float('inf'):
            speedup = cpu_time / gpu_time
        else:
            speedup = float('nan')
        
        print(f"{size:8d}  |  {gpu_time:12.3f}  |  {cpu_time:12.3f}  |  {speedup:12.1f}x")

if __name__ == "__main__":
    main() 